/**
 * test_stream_pusher.cpp
 * 
 * 测试 stream_pusher.h 的推流功能
 * 支持读取 h264/h265/mp4 等格式并推流到 RTSP/RTMP/文件
 */

#include "stream_pusher.h"
#include "stream_demuxer.h"
#include <iostream>
#include <string>
#include <chrono>
#include <thread>
#include <future>
#include <signal.h>

using namespace enrigin;

// 全局标志用于优雅退出
static bool g_running = true;

void signal_handler(int sig) {
    std::cout << "\nReceived signal " << sig << ", stopping..." << std::endl;
    g_running = false;
}

// 简单的观察者实现用于同步读取
class SimpleDemuxerObserver : public StreamDemuxerEvents {
public:
    AVFormatContext* ifmt_ctx = nullptr;
    FfmpegOutputer* outputer = nullptr;
    
    // 统计信息
    std::atomic<int> packet_count{0};
    std::atomic<int> video_packet_count{0};
    std::atomic<int> audio_packet_count{0};
    std::atomic<int> error_count{0};
    bool first_keyframe_sent = false;
    
    // 时间控制
    std::chrono::steady_clock::time_point start_time;
    std::chrono::steady_clock::time_point last_report_time;
    
    void onAvformatOpened(AVFormatContext *iFmtCtx) override {
        ifmt_ctx = iFmtCtx;
        start_time = std::chrono::steady_clock::now();
        last_report_time = start_time;
    }
    
    void onAvformatClosed() override {
        ifmt_ctx = nullptr;
    }
    
    void onReadEof(AVPacket *pkt) override {
        std::cout << "\nReached end of file" << std::endl;
        g_running = false;
    }
    
    int onReadFrame(AVPacket *packet) override {
        if (!packet || !ifmt_ctx || !outputer) {
            return -1;
        }
        
        packet_count++;
        
        // 统计音视频包
        if (packet->stream_index >= 0 && packet->stream_index < (int)ifmt_ctx->nb_streams) {
            AVStream* stream = ifmt_ctx->streams[packet->stream_index];
            if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
                video_packet_count++;
            } else if (stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
                audio_packet_count++;
            }
        }
        
        // 跳过音频（暂时只推视频）
        if (packet->stream_index >= 0 && packet->stream_index < (int)ifmt_ctx->nb_streams) {
            AVStream* stream = ifmt_ctx->streams[packet->stream_index];
            
            // 只处理视频流
            if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
                // 基于时间戳的速率控制（模拟实时推流）
                static bool rate_control_enabled = false;
                static int rate_control_count = 0;
                if (packet->pts != AV_NOPTS_VALUE) {
                    if (!rate_control_enabled) {
                        std::cout << "\n[RATE_CONTROL] Enabled for video stream" << std::endl;
                        rate_control_enabled = true;
                    }
                    
                    double pts_time = packet->pts * av_q2d(stream->time_base);
                    auto target_time = start_time + std::chrono::duration_cast<std::chrono::steady_clock::duration>(
                        std::chrono::duration<double>(pts_time));
                    
                    auto now = std::chrono::steady_clock::now();
                    auto sleep_duration = target_time - now;
                    if (sleep_duration.count() > 0) {
                        if (rate_control_count++ < 3) {
                            std::cout << "[RC] Sleep " << std::chrono::duration_cast<std::chrono::milliseconds>(sleep_duration).count() << "ms" << std::endl;
                        }
                        std::this_thread::sleep_for(sleep_duration);
                    }
                }
                
                // 推送到输出流
                int ret = outputer->inputPacket(packet);
                if (ret != 0) {
                    error_count++;
                    if (error_count < 10) {
                        std::cerr << "\nFailed to input packet (error " << error_count << ")" << std::endl;
                    }
                }
            }
        }
        
        // 定期报告进度
        auto now = std::chrono::steady_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(now - last_report_time).count();
        if (elapsed >= 1) {
            auto total_elapsed = std::chrono::duration_cast<std::chrono::seconds>(now - start_time).count();
            std::cout << "\r[" << total_elapsed << "s] Packets: " << packet_count 
                      << " (Video: " << video_packet_count 
                      << ", Audio: " << audio_packet_count << ")" << std::flush;
            last_report_time = now;
        }
        
        return 0;
    }
    
};

void print_usage(const char* prog_name) {
    std::cout << "Usage: " << prog_name << " <input_file> <output_url> [preset]" << std::endl;
    std::cout << std::endl;
    std::cout << "Arguments:" << std::endl;
    std::cout << "  input_file   : Input video file (h264/h265/mp4/mkv/etc.)" << std::endl;
    std::cout << "  output_url   : Output URL or file path" << std::endl;
    std::cout << "                 - RTSP: rtsp://host:port/path" << std::endl;
    std::cout << "                 - RTMP: rtmp://host:port/path" << std::endl;
    std::cout << "                 - UDP:  udp://host:port" << std::endl;
    std::cout << "                 - File: output.mp4, output.h264, etc." << std::endl;
    std::cout << "  preset       : (Optional) Timestamp smoothing preset" << std::endl;
    std::cout << "                 - default: balanced settings" << std::endl;
    std::cout << "                 - conservative: minimal intervention" << std::endl;
    std::cout << "                 - aggressive: strong smoothing" << std::endl;
    std::cout << "                 - looping: for file looping scenarios" << std::endl;
    std::cout << std::endl;
    std::cout << "Examples:" << std::endl;
    std::cout << "  " << prog_name << " input.h264 rtsp://10.41.216.1:8554/live/test" << std::endl;
    std::cout << "  " << prog_name << " input.mp4 rtmp://10.41.216.1/live/stream" << std::endl;
    std::cout << "  " << prog_name << " input.h265 output.mp4 aggressive" << std::endl;
    std::cout << "  " << prog_name << " input.mp4 udp://127.0.0.1:5000" << std::endl;
}

int main(int argc, char* argv[]) {
    // 参数检查
    if (argc < 3) {
        print_usage(argv[0]);
        return 1;
    }

    std::string input_file = argv[1];
    std::string output_url = argv[2];
    std::string preset = (argc >= 4) ? argv[3] : "default";

    std::cout << "========================================" << std::endl;
    std::cout << "Stream Pusher Test" << std::endl;
    std::cout << "========================================" << std::endl;
    std::cout << "Input:  " << input_file << std::endl;
    std::cout << "Output: " << output_url << std::endl;
    std::cout << "Preset: " << preset << std::endl;
    std::cout << "========================================" << std::endl;

    // 注册信号处理
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    // 1. 创建解复用器读取输入文件
    SimpleDemuxerObserver observer;
    StreamDemuxer demuxer;
    
    // 使用同步模式打开流（isSyncOpen=true）
    int ret = demuxer.openStream(input_file, &observer, false, true);
    if (ret != 0) {
        std::cerr << "Failed to open input file: " << input_file << std::endl;
        return 1;
    }

    std::cout << "Input stream opened successfully" << std::endl;
    
    // 获取输入格式上下文
    AVFormatContext* ifmt_ctx = observer.ifmt_ctx;
    if (!ifmt_ctx || ifmt_ctx->nb_streams == 0) {
        std::cerr << "Invalid input format context" << std::endl;
        demuxer.closeStream(false);
        return 1;
    }

    // 打印输入流信息
    std::cout << "Input stream info:" << std::endl;
    std::cout << "  Duration: " << ifmt_ctx->duration / AV_TIME_BASE << " seconds" << std::endl;
    std::cout << "  Streams: " << ifmt_ctx->nb_streams << std::endl;
    for (unsigned int i = 0; i < ifmt_ctx->nb_streams; i++) {
        AVStream* stream = ifmt_ctx->streams[i];
        AVCodecParameters* codecpar = stream->codecpar;
        std::cout << "  Stream #" << i << ": ";
        if (codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            std::cout << "Video - " << avcodec_get_name(codecpar->codec_id)
                      << " " << codecpar->width << "x" << codecpar->height;
        } else if (codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
            std::cout << "Audio - " << avcodec_get_name(codecpar->codec_id);
        } else {
            std::cout << "Other";
        }
        std::cout << std::endl;
    }

    // 2. 创建推流器
    FfmpegOutputer outputer;
    
    // 设置时间戳平滑预设
    // 对于RTMP推流，使用更激进的平滑策略
    if (output_url.find("rtmp://") != std::string::npos) {
        std::cout << "RTMP detected, using aggressive timestamp smoothing" << std::endl;
        outputer.setTimestampSmoothingPreset("aggressive");
    } else {
        outputer.setTimestampSmoothingPreset(preset);
    }
    
    // 打开输出流
    ret = outputer.openOutputStream(output_url, ifmt_ctx);
    if (ret != 0) {
        std::cerr << "Failed to open output stream: " << output_url << std::endl;
        demuxer.closeStream(false);
        return 1;
    }

    std::cout << "Output stream opened successfully" << std::endl;
    
    // 3. 设置回调的outputer并开始自动推流
    observer.outputer = &outputer;
    
    std::cout << "Starting to push packets..." << std::endl;
    std::cout << "(Press Ctrl+C to stop)" << std::endl;
    std::cout << "========================================" << std::endl;

    // StreamDemuxer会自动通过onReadFrame回调传递packet
    // 等待处理完成或用户中断
    while (g_running) {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    
    // 等待输出队列清空
    std::cout << "\nFlushing output queue..." << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(2));

    std::cout << std::endl;
    std::cout << "========================================" << std::endl;
    std::cout << "Streaming completed" << std::endl;
    std::cout << "Total packets: " << observer.packet_count << std::endl;
    std::cout << "  Video packets: " << observer.video_packet_count << std::endl;
    std::cout << "  Audio packets: " << observer.audio_packet_count << std::endl;
    std::cout << "  Error count: " << observer.error_count << std::endl;
    std::cout << "  First keyframe sent: " << (observer.first_keyframe_sent ? "Yes" : "No") << std::endl;

    // 获取时间戳平滑统计
    int64_t total_pkts, corrected_pkts;
    double correction_rate;
    outputer.getTimestampStatistics(total_pkts, corrected_pkts, correction_rate);
    std::cout << "Timestamp smoothing:" << std::endl;
    std::cout << "  Total packets: " << total_pkts << std::endl;
    std::cout << "  Corrected packets: " << corrected_pkts << std::endl;
    std::cout << "  Correction rate: " << (correction_rate * 100.0) << "%" << std::endl;

    auto total_time = std::chrono::duration_cast<std::chrono::seconds>(
        std::chrono::steady_clock::now() - observer.start_time).count();
    std::cout << "Total time: " << total_time << " seconds" << std::endl;
    std::cout << "========================================" << std::endl;

    // 4. 清理资源
    std::cout << "\nCleaning up resources..." << std::endl;
    
    // 先关闭输出流，确保所有数据都写入
    std::cout << "Closing output stream..." << std::endl;
    outputer.closeOutputStream();
    
    // 再关闭输入流
    std::cout << "Closing input stream..." << std::endl;
    demuxer.closeStream(false);

    std::cout << "Test completed successfully" << std::endl;
    
    // 如果有错误，返回非零退出码
    if (observer.error_count > 0) {
        std::cout << "Warning: " << observer.error_count << " errors occurred during streaming" << std::endl;
        return 2;
    }
    
    return 0;
}
